from pathlib import Path
import torch
from ultralytics import YOLO

project_dir = Path(__file__).parent
data_yaml = (project_dir / 'data' / 'hand-keypoints' / 'data.yaml').absolute()
device = 'cuda' if torch.cuda.is_available() else 'cpu'

assert data_yaml.exists()

print(f'device: {device}')

model = YOLO('yolo11n-pose.pt')
model.train(data=data_yaml, epochs=100, imgsz=640, batch=8, device=device)